import numpy as np
import multiprocessing as mp
#from collections import defaultdict
#from itertools import chain
import os
import time
import pickle


def save_new_dict(i):

    point_dict_grasp = np.load('mask_label_grasp/image_{:05d}.npz'.format(i),allow_pickle=True)['arr_0'].item()
    point_dict_suction = np.load('mask_label_suction/image_{:05d}.npz'.format(i),allow_pickle=True)['arr_0'].item()

    all = defaultdict(list)

    for k, v in chain(point_dict_grasp.items(), point_dict_suction.items()):
        all[k].extend(v)

    save_path = os.path.join('mask_label_all','image_{:05d}.pickle'.format(i))

    with open(save_path,'wb') as f:
         pickle.dump(all, f, protocol = pickle.HIGHEST_PROTOCOL)

    print('{} is finished'.format(i))



def save_new_npy(i):
    label_dict = {}
    time_start = time.time()
    mask_label_grasp = np.load('mask_label_grasp/image_{:05d}.npy'.format(i))
    mask_label_suction = np.load('mask_label_suction/image_{:05d}.npy'.format(i))
    all = np.concatenate((mask_label_grasp, mask_label_suction), axis=1)
    data = all.astype(np.float16)
    save_path = os.path.join('mask_label_all/', 'image_{:05d}.npy'.format(i))
    np.save(save_path, data)

    point_label = np.load('./mask_label_all/image_{:05d}.npy'.format(i))
    camera_intr = np.load('./mask_label_all/camera_intr/camera_intr_{:05d}.npy'.format(i))
    camera_pose = np.load('./mask_label_all/camera_pose/camera_pose_{:05d}.npy'.format(i))
    normal = np.load('./mask_label_all/pointcloud_data_normal/xyz_normal_{:05d}.npy'.format(i))
    depth_img = np.load('./mask_label_all/depth/depth_{:05d}.npy'.format(i))
    label_dict['point_label'] = point_label
    label_dict['camera_intr'] = camera_intr
    label_dict['camera_pose'] = camera_pose
    label_dict['point_normal'] = normal
    label_dict['depth_img'] = depth_img
    #print(point_label.shape)
    #print(camera_intr.shape, camera_pose.shape, normal.shape, depth_img.shape)
    #with open('./mask_label_all/image_{:05d}.pickle'.format(i), 'wb') as fp:
        #pickle.dump(label_dict, fp, protocol = pickle.HIGHEST_PROTOCOL)
    np.save('./mask_label_all/label_{:05d}.npy'.format(i), label_dict)
    #print('time_cost is :', time.time() - time_start)
        


    print('{} is finished'.format(i))

if __name__ == "__main__":
    use_mp = True

    if use_mp:
        pool = mp.Pool(processes=mp.cpu_count()*4)

    for i in range(0, 1):

        if use_mp:
            pool.apply_async(save_new_npy, args=(i,))
        else:
            print('single process for debug')
            save_new_npy(i)

    if use_mp:
        pool.close()
        pool.join()





